package lc.Q200_250;

import lc.TreeNode;

import java.util.HashMap;
import java.util.Map;

/**
 * 给定一个二叉搜索树的根节点 root ，和一个整数 k ，请你设计一个算法查找其中第 k 小的元素（从 1 开始计数）。
 */
public class Q230 {
    Map<TreeNode, Integer> counter = new HashMap<TreeNode, Integer>();
    public int kthSmallest(TreeNode root, int k) {
        // 二叉搜索树，左边小，右边大
        // 数下两边个数
        countNodes(root);
        int leftCount = counter.getOrDefault(root.left, 0);
        if (leftCount == k - 1) {
            // 说明就是根节点
            return root.val;
        }
        if (leftCount >= k) {
            // 在左边
            return kthSmallest(root.left, k);
        }
        return kthSmallest(root.right, k - 1 - leftCount);
    }

    public int countNodes(TreeNode root) {
        if (root == null) {
            return 0;
        }
        int i = 1 + countNodes(root.left) + countNodes(root.right);
        counter.put(root, i);
        return i;
    }

    public static void main(String[] args) {
        TreeNode root = new TreeNode(5,
                new TreeNode(3, new TreeNode(
                        2, new TreeNode(1, null, null), null
                ), new TreeNode(4)),
                new TreeNode(6, null, null)
                );
        Q230 q230 = new Q230();
        System.out.println(q230.kthSmallest(root, 3));
    }
}
